import torch
import pandas as pd
import random
import time

filename = '美的集团000333.SZ.xlsx'
# filename='中国平安601318.SH.xlsx'
# filename = '分众传媒002027.SZ.xlsx'
# filename = '立讯精密002475.SZ.xlsx'
file_path = r'C:\Users\masaikk\Desktop\data_aug_2\{}'.format(filename)
raw_file_data = pd.read_excel(file_path)
file_data = raw_file_data.head(4000)
init_positions = file_data.iloc[0].收盘价
init_positions_number = 1500000 / init_positions
# init_money = 1033240
init_money = 1000000
init_current_line = 0
action_choose = [0, 1, 2, 3, 4, 5, 6, 7, 8]
model = torch.nn.Sequential(
    torch.nn.Linear(12, 128),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 9),
)
model.load_state_dict(
    torch.load(
        r'C:\Users\masaikk\Downloads\100model_128_1668085603_500_美的集团000333.SZ.xlsx_-166008521.713676.pth'))


class StockEnv:
    def __init__(self):
        self.stock_data = file_data
        self.over_line = len(self.stock_data)
        self.money = init_money
        self.action_choose = action_choose
        self.current_line = init_current_line
        self.positions_number = init_positions_number
        self.positions_price = init_positions
        self.rating = 2  # 赚钱的比例系数

    def step(self, action):
        this_time_stock = self.stock_data.iloc[self.current_line]
        action2reward = -10
        if action == 0:
            # 买100手
            self.positions_price = ((self.positions_price * self.positions_number) + 10000 *
                                    this_time_stock['价']) / (self.positions_number + 10000)
            self.positions_number = self.positions_number + 10000
            self.money = self.money - 10000 * this_time_stock['价'] - 0.00015 * 10000 * this_time_stock['价']
            action2reward = -200
            if self.money < 0:
                action2reward = action2reward * 1000000
        elif action == 1:
            # 买50手
            self.positions_price = ((self.positions_price * self.positions_number) + 5000 *
                                    this_time_stock['价']) / (self.positions_number + 5000)
            self.positions_number = self.positions_number + 5000
            self.money = self.money - 5000 * this_time_stock['价'] - 0.00015 * 5000 * this_time_stock['价']
            action2reward = -100

            if self.money < 0:
                action2reward = action2reward * 1000000
        elif action == 2:
            # 买20手
            self.positions_price = ((self.positions_price * self.positions_number) + 2000 *
                                    this_time_stock['价']) / (self.positions_number + 2000)
            self.positions_number = self.positions_number + 2000
            self.money = self.money - 2000 * this_time_stock['价'] - 0.00015 * 2000 * this_time_stock['价']
            action2reward = -40
            if self.money < 0:
                action2reward = action2reward * 1000000
        elif action == 3:
            # 买10手
            self.positions_price = ((self.positions_price * self.positions_number) + 1000 *
                                    this_time_stock['价']) / (self.positions_number + 1000)
            self.positions_number = self.positions_number + 1000
            self.money = self.money - 1000 * this_time_stock['价'] - 0.00015 * 1000 * this_time_stock['价']
            action2reward = -20
            if self.money < 0:
                action2reward = action2reward * 1000000
        elif action == 4:
            pass
        elif action == 5:

            if self.positions_number < 1000:
                action2reward = -10000000
            else:
                self.positions_number = self.positions_number - 1000
                self.money = self.money + 1000 * this_time_stock['价'] - 0.00115 * 1000 * this_time_stock['价']
                # 卖10手
                action2reward = self.rating * (
                        this_time_stock['价'] - self.positions_price) * 1000 - 10  # 赚到0元也要扣10reward
                self.positions_number = self.positions_number - 1000
        elif action == 6:

            if self.positions_number < 2000:
                action2reward = -20000000
            else:
                self.positions_number = self.positions_number - 2000
                self.money = self.money + 2000 * this_time_stock['价'] - 0.00115 * 2000 * this_time_stock['价']
                # 卖20手
                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 20 - 10
                self.positions_number = self.positions_number - 2000
        elif action == 7:

            if self.positions_number < 5000:
                action2reward = -50000000
            else:
                self.positions_number = self.positions_number - 5000
                self.money = self.money + 5000 * this_time_stock['价'] - 0.00115 * 5000 * this_time_stock['价']
                # 卖50手
                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 5000 - 10
                self.positions_number = self.positions_number - 5000
        elif action == 8:

            if self.positions_number < 10000:
                action2reward = -100000000
            else:
                self.positions_number = self.positions_number - 10000
                self.money = self.money + 10000 * this_time_stock['价'] - 0.00115 * 10000 * this_time_stock['价']
                # 卖100手
                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 10000 - 10
                self.positions_number = self.positions_number - 10000
        else:
            pass

        self.current_line += 1
        current_stock = self.stock_data.iloc[self.current_line]
        next_state = (current_stock['收盘价'], current_stock['开盘价'], current_stock['最高价'],
                      current_stock['最低价'], current_stock['成交额（元）'], current_stock['成交量（手）'],
                      current_stock['当前额'], current_stock['当前量'], current_stock['价'],
                      self.positions_number, self.positions_price, self.money)
        over = self.current_line == self.over_line - 1
        return next_state, action2reward, over, None

    def get_current_state(self):
        current_stock = self.stock_data.iloc[self.current_line]
        state = (current_stock['收盘价'], current_stock['开盘价'], current_stock['最高价'],
                 current_stock['最低价'], current_stock['成交额（元）'], current_stock['成交量（手）'],
                 current_stock['当前额'], current_stock['当前量'], current_stock['价'],
                 self.positions_number, self.positions_price, self.money)
        return state

    def get_action(self, state):
        if random.random() < 0.01:
            return random.choice(self.action_choose)

        # 走神经网络,得到一个动作
        # print('走神经网络,得到一个动作')
        reshape_state = torch.FloatTensor(state).reshape(1, 12)

        return model(reshape_state).argmax().item()

    def reset(self):
        self.current_line = init_current_line
        self.positions_number = init_positions_number
        self.positions_price = init_positions

        return self.get_current_state()


env = StockEnv()

model.eval()
env.reset()
test_state = env.get_current_state()
reshape_state = torch.FloatTensor(test_state).reshape(1, 12)
log_data = []
for turn in range(10000):
    action = model(reshape_state).argmax().item()
    # print(action)
    next_state, action2reward, over, _ = env.step(action)
    if over:
        env.reset()
    print("action {}, reward {}".format(action, action2reward))
    log_data.append(
        [action, next_state[11], next_state[9], next_state[10], next_state[11] + next_state[9] * next_state[10]])
    reshape_state = torch.FloatTensor(next_state).reshape(1, 12)
log_data_csv = pd.DataFrame(log_data, columns=['action', 'now_money', 'positions_number', 'positions_price', 'total'])
log_data_csv.to_csv('./logs/{}_{}.csv'.format(int(time.time()), filename))
